package ixRedis

import (
	"context"
	"time"

	"github.com/garyburd/redigo/redis"
)

// RateLimiter 限流器
type RateLimiter struct {
	client *Client
	key    string
	limit  int64
	window time.Duration
}

// NewRateLimiter 创建限流器
func (c *Client) NewRateLimiter(key string, limit int64, window time.Duration) *RateLimiter {
	return &RateLimiter{
		client: c,
		key:    "rate_limit:" + key,
		limit:  limit,
		window: window,
	}
}

// Allow 检查是否允许请求
func (rl *RateLimiter) Allow(ctx context.Context) (bool, error) {
	conn, err := rl.client.getConn()
	if err != nil {
		return false, err
	}
	defer conn.Close()

	now := time.Now().Unix()

	// 使用 Lua 脚本实现滑动窗口限流
	script := `
		local key = KEYS[1]
		local limit = tonumber(ARGV[1])
		local window = tonumber(ARGV[2])
		local now = tonumber(ARGV[3])
		
		-- 清理过期的记录
		redis.call('ZREMRANGEBYSCORE', key, '-inf', now - window)
		
		-- 获取当前窗口内的请求数
		local current = redis.call('ZCARD', key)
		
		if current < limit then
			-- 添加当前请求
			redis.call('ZADD', key, now, now)
			redis.call('EXPIRE', key, window)
			return {1, limit - current - 1}
		else
			return {0, 0}
		end
	`

	result, err := redis.Ints(conn.Do("EVAL", script, 1, rl.key, rl.limit, int64(rl.window.Seconds()), now))
	if err != nil {
		return false, err
	}

	allowed := result[0] == 1

	return allowed, nil
}

// AllowN 检查是否允许 N 个请求
func (rl *RateLimiter) AllowN(ctx context.Context, n int64) (bool, int64, error) {
	conn, err := rl.client.getConn()
	if err != nil {
		return false, 0, err
	}
	defer conn.Close()

	now := time.Now().Unix()

	// 使用 Lua 脚本实现滑动窗口限流
	script := `
		local key = KEYS[1]
		local limit = tonumber(ARGV[1])
		local window = tonumber(ARGV[2])
		local now = tonumber(ARGV[3])
		local n = tonumber(ARGV[4])
		
		-- 清理过期的记录
		redis.call('ZREMRANGEBYSCORE', key, '-inf', now - window)
		
		-- 获取当前窗口内的请求数
		local current = redis.call('ZCARD', key)
		
		if current + n <= limit then
			-- 添加 N 个请求
			for i = 1, n do
				redis.call('ZADD', key, now + i, now + i)
			end
			redis.call('EXPIRE', key, window)
			return {1, limit - current - n}
		else
			return {0, 0}
		end
	`

	result, err := redis.Ints(conn.Do("EVAL", script, 1, rl.key, rl.limit, int64(rl.window.Seconds()), now, n))
	if err != nil {
		return false, 0, err
	}

	allowed := result[0] == 1
	remaining := int64(result[1])

	return allowed, remaining, nil
}

// GetRemaining 获取剩余请求数
func (rl *RateLimiter) GetRemaining(ctx context.Context) (int64, error) {
	conn, err := rl.client.getConn()
	if err != nil {
		return 0, err
	}
	defer conn.Close()

	now := time.Now().Unix()
	windowStart := now - int64(rl.window.Seconds())

	// 清理过期的记录
	conn.Do("ZREMRANGEBYSCORE", rl.key, "-inf", windowStart)

	// 获取当前窗口内的请求数
	current, err := redis.Int64(conn.Do("ZCARD", rl.key))
	if err != nil {
		return 0, err
	}

	remaining := rl.limit - current
	if remaining < 0 {
		remaining = 0
	}

	return remaining, nil
}

// Reset 重置限流器
func (rl *RateLimiter) Reset(ctx context.Context) error {
	_, err := rl.client.Del(rl.key)
	return err
}

// TokenBucketRateLimiter 令牌桶限流器
type TokenBucketRateLimiter struct {
	client     *Client
	key        string
	capacity   int64
	refillRate int64
	refillTime time.Duration
}

// NewTokenBucketRateLimiter 创建令牌桶限流器
func (c *Client) NewTokenBucketRateLimiter(key string, capacity, refillRate int64, refillTime time.Duration) *TokenBucketRateLimiter {
	return &TokenBucketRateLimiter{
		client:     c,
		key:        "token_bucket:" + key,
		capacity:   capacity,
		refillRate: refillRate,
		refillTime: refillTime,
	}
}

// Allow 检查是否允许请求
func (tbrl *TokenBucketRateLimiter) Allow(ctx context.Context) (bool, error) {
	conn, err := tbrl.client.getConn()
	if err != nil {
		return false, err
	}
	defer conn.Close()

	now := time.Now().Unix()

	// 使用 Lua 脚本实现令牌桶算法
	script := `
		local key = KEYS[1]
		local capacity = tonumber(ARGV[1])
		local refillRate = tonumber(ARGV[2])
		local refillTime = tonumber(ARGV[3])
		local now = tonumber(ARGV[4])
		
		local bucket = redis.call('HMGET', key, 'tokens', 'lastRefill')
		local tokens = tonumber(bucket[1]) or capacity
		local lastRefill = tonumber(bucket[2]) or now
		
		-- 计算需要添加的令牌数
		local timePassed = now - lastRefill
		local tokensToAdd = math.floor(timePassed / refillTime) * refillRate
		
		-- 更新令牌数
		tokens = math.min(capacity, tokens + tokensToAdd)
		lastRefill = now
		
		if tokens >= 1 then
			tokens = tokens - 1
			redis.call('HMSET', key, 'tokens', tokens, 'lastRefill', lastRefill)
			redis.call('EXPIRE', key, 3600) -- 1小时过期
			return {1, tokens}
		else
			redis.call('HMSET', key, 'tokens', tokens, 'lastRefill', lastRefill)
			redis.call('EXPIRE', key, 3600) -- 1小时过期
			return {0, tokens}
		end
	`

	result, err := redis.Ints(conn.Do("EVAL", script, 1, tbrl.key, tbrl.capacity, tbrl.refillRate, int64(tbrl.refillTime.Seconds()), now))
	if err != nil {
		return false, err
	}

	allowed := result[0] == 1

	return allowed, nil
}

// GetTokens 获取当前令牌数
func (tbrl *TokenBucketRateLimiter) GetTokens(ctx context.Context) (int64, error) {
	conn, err := tbrl.client.getConn()
	if err != nil {
		return 0, err
	}
	defer conn.Close()

	now := time.Now().Unix()

	// 使用 Lua 脚本获取当前令牌数
	script := `
		local key = KEYS[1]
		local capacity = tonumber(ARGV[1])
		local refillRate = tonumber(ARGV[2])
		local refillTime = tonumber(ARGV[3])
		local now = tonumber(ARGV[4])
		
		local bucket = redis.call('HMGET', key, 'tokens', 'lastRefill')
		local tokens = tonumber(bucket[1]) or capacity
		local lastRefill = tonumber(bucket[2]) or now
		
		-- 计算需要添加的令牌数
		local timePassed = now - lastRefill
		local tokensToAdd = math.floor(timePassed / refillTime) * refillRate
		
		-- 更新令牌数
		tokens = math.min(capacity, tokens + tokensToAdd)
		lastRefill = now
		
		redis.call('HMSET', key, 'tokens', tokens, 'lastRefill', lastRefill)
		redis.call('EXPIRE', key, 3600) -- 1小时过期
		
		return tokens
	`

	tokens, err := redis.Int64(conn.Do("EVAL", script, 1, tbrl.key, tbrl.capacity, tbrl.refillRate, int64(tbrl.refillTime.Seconds()), now))
	if err != nil {
		return 0, err
	}

	return tokens, nil
}

// Reset 重置令牌桶
func (tbrl *TokenBucketRateLimiter) Reset(ctx context.Context) error {
	_, err := tbrl.client.Del(tbrl.key)
	return err
}
